import numpy as np
import os
import struct

MAX_STRING_LEN = 256
HF_MAX_STR_LEN = 128
QT_MAX_STR_LEN = 64


class NameLoc:
    def __init__(self):
        self.type = "u"
        self.length = 0
        self.on_grid = False
        self.time_type = ""
        self.time = 0.0
        self.Max = 0.0
        self.Min = 0.0
        self.Avg = 0.0
        self.name = ""
        self.unit = ""
        self.fcount = 0
        self.offset = 0


class WellData:
    def __init__(self):
        self.XCoord = 0.0
        self.YCoord = 0.0
        self.timesteps = 0
        self.well_num = 0
        self.points_num = 0
        self.stat_data = []
        self.dym_data = []


class BinReader:
    def __init__(self):
        self.clear()

    def clear(self):
        self.casename = ""
        self.stat_props = []  # Static property list
        self.dym_props = []  # Dynamic property list
        self.well_data = WellData()  # Well data list
        self.Gross2Net = np.array([], dtype=np.int32)
        self.Net2Gross = np.array([], dtype=np.int32)
        self.GridInfo = np.zeros(32, dtype=np.int32)
        self.AxisOffset = np.zeros(3, dtype=np.float64)
        self.origin_net_size = 0
        self.origin_GridInfo = np.zeros(32, dtype=np.int32)
        self.origin_Gross2Net = np.array([], dtype=np.int32)

    def print_info(self):
        # show class real length
        print("case name size: ", len(self.casename))
        print("stat_props size: ", len(self.stat_props))
        print("dym_props size: ", len(self.dym_props))
        print("well_data size: ", len(self.well_data.stat_data))
        print("Gross2Net size: ", len(self.Gross2Net))
        print("Net2Gross size: ", len(self.Net2Gross))
        print("GridInfo size: ", len(self.GridInfo))
        print("AxisOffset size: ", len(self.AxisOffset))
        # print("origin_net_size size: ", len(self.origin_net_size))
        print("origin_GridInfo size: ", len(self.origin_GridInfo))
        print("origin_Gross2Net size: ", len(self.origin_Gross2Net))

    # 0-Cartesian, 1-Stage, 2-CPG, 3-PEBI
    def get_grid_type(self):
        if self.GridInfo[0] == 0:
            return "Cartesian"
        elif self.GridInfo[0] == 1:
            return "Stage"
        elif self.GridInfo[0] == 2:
            return "CPG"
        elif self.GridInfo[0] == 3:
            return "PEBI"

    def get_porolvl(self):
        return self.GridInfo[1]  # 双孔模型层数

    def get_grid_ni(self):
        return self.GridInfo[3]

    def get_grid_nj(self):
        return self.GridInfo[4]

    def get_grid_nk(self):
        return self.GridInfo[5]

    def get_net_size(self):
        return len(self.Net2Gross)  # GridInfo[8]

    def get_gross_size(self):
        return len(self.Gross2Net)  # GridInfo[2]

    def get_time_steps(self):
        return len(self.dym_props)

    def get_base_grid_num(self):
        return self.GridInfo[2] - self.GridInfo[6] - self.GridInfo[7]

    def get_active_base_grid_num(self):
        return self.GridInfo[9] + self.GridInfo[11] * (self.GridInfo[1] - 1)

    def get_fine_grid_num(self):
        return self.GridInfo[6]

    def get_active_fine_grid_num(self):
        return self.GridInfo[10] + self.GridInfo[12] * (self.GridInfo[1] - 1)

    def scan_bin_array(self, fp, a):
        """
        Read a binary array from the file pointer.

        :param fp: File pointer.
        :param a: NameLoc object.

        """
        a.offset = fp.tell()

        # Read char_len (UGINT64, 8 bytes, unsigned long long)
        char_len, a.time = struct.unpack("<q d", fp.read(16))
        a.type = fp.read(1).decode("latin1")

        # Read pname (HF_MAX_STR_LEN bytes)
        a.name = fp.read(HF_MAX_STR_LEN).split(b"\x00", 1)[0].decode("latin1")

        # Read punit (HF_MAX_STR_LEN bytes)
        punit_bytes = fp.read(HF_MAX_STR_LEN)
        a.unit = punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")

        offset = QT_MAX_STR_LEN
        a.Max, a.Min, a.Avg = struct.unpack("<ddd", punit_bytes[offset : offset + 24])
        offset += 24
        a.time_type = punit_bytes[offset : offset + 1].decode("latin1")
        a.on_grid = punit_bytes[offset + 1 : offset + 2] == b"y"

        net_len = char_len - 8 - 1 - MAX_STRING_LEN
        fp.seek(net_len, 1)

        if a.type == "d":
            a.length = net_len // 8
        elif a.type == "i":
            a.length = net_len // 4
        else:
            raise ValueError(f"Unknown data type: {a.type}")

    def get_static_data(self, name, prop_array, remap=False):
        """
        Get static data from the binary file.

        :param name: Property name.
        :param prop_array: Dictionary to store the property data.
        :param remap: Whether to remap the property data.

        :return: True if the property is found, False otherwise.

        """
        prop = next((p for p in self.stat_props if p.name == name), None)
        if not prop:
            return False

        filename = f"{self.casename}_geom.bin"
        try:
            with open(filename, "rb") as stat_prop_fp:
                # Read char_len (UGINT64, 8 bytes, unsigned long long)
                stat_prop_fp.seek(prop.offset, os.SEEK_SET)
                char_len, prop_array["time"] = struct.unpack(
                    "<q d", stat_prop_fp.read(16)
                )
                prop_array["type"] = stat_prop_fp.read(1).decode("latin1")

                # Read pname (HF_MAX_STR_LEN bytes)
                prop_array["name"] = (
                    stat_prop_fp.read(HF_MAX_STR_LEN)
                    .split(b"\x00", 1)[0]
                    .decode("latin1")
                )

                # Read punit (HF_MAX_STR_LEN bytes)
                punit_bytes = stat_prop_fp.read(HF_MAX_STR_LEN)
                prop_array["unit"] = (
                    punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")
                )

                # Extract Max, Min, Avg, time_type, and on_grid
                offset = QT_MAX_STR_LEN
                prop_array["Max"], prop_array["Min"], prop_array["Avg"] = struct.unpack(
                    "<ddd", punit_bytes[offset : offset + 24]
                )
                prop_array["time_type"] = punit_bytes[offset + 24 : offset + 25].decode(
                    "latin1"
                )
                prop_array["on_grid"] = punit_bytes[offset + 25 : offset + 26] == b"y"

                # Calculate net_len and read data
                net_len = char_len - 8 - 1 - MAX_STRING_LEN
                if prop_array["type"] == "d":
                    prop_array["length"] = net_len // 8
                    prop_array["data"] = np.frombuffer(
                        stat_prop_fp.read(net_len), dtype=np.float64
                    )
                elif prop_array["type"] == "i":
                    prop_array["length"] = net_len // 4
                    prop_array["data"] = np.frombuffer(
                        stat_prop_fp.read(net_len), dtype=np.int32
                    )
                else:
                    raise ValueError(f"Unknown data type: {prop_array['type']}")

                if np.isnan(prop_array["data"]).any():
                    print(f"Warning: NaN values found in property {name}")
                if np.isinf(prop_array["data"]).any():
                    print(f"Warning: Inf values found in property {name}")

                # Handle remapping if needed
                if remap and prop_array["on_grid"]:
                    if hasattr(self, "Moves") and self.Moves:
                        prop_array["data"] = self.re_map(
                            prop_array["data"],
                            self.NetSize(),
                            self.origin_net_size,
                            self.Moves,
                        )

                return True
        except FileNotFoundError:
            print(f"Error: File '{filename}' not found.")
            return False

    def get_dynamic_data(self, tstep, name, prop_array, remap=False):
        """
        Get dynamic data from the binary file.

        :param tstep: Time step.
        :param name: Property name.
        :param prop_array: Dictionary to store the property data.
        :param remap: Whether to remap the property data.

        :return: True if the property is found, False otherwise.

        """
        if tstep > len(self.dym_props):
            return False

        props = self.dym_props[tstep - 1]
        prop = next((p for p in props if p.name == name), None)
        if not prop:
            return False

        filename = f"{self.casename}_stu{prop.fcount if prop.fcount > 0 else ''}.bin"
        try:
            with open(filename, "rb") as dym_prop_fp:
                dym_prop_fp.seek(prop.offset, os.SEEK_SET)
                char_len, prop_array["time"] = struct.unpack(
                    "<q d", dym_prop_fp.read(16)
                )
                prop_array["type"] = dym_prop_fp.read(1).decode("latin1")

                # Read pname (HF_MAX_STR_LEN bytes)
                prop_array["name"] = (
                    dym_prop_fp.read(HF_MAX_STR_LEN)
                    .split(b"\x00", 1)[0]
                    .decode("latin1")
                )

                # Read punit (HF_MAX_STR_LEN bytes)
                punit_bytes = dym_prop_fp.read(HF_MAX_STR_LEN)
                prop_array["unit"] = (
                    punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")
                )

                # Extract Max, Min, Avg, time_type, and on_grid
                offset = QT_MAX_STR_LEN
                prop_array["Max"], prop_array["Min"], prop_array["Avg"] = struct.unpack(
                    "<ddd", punit_bytes[offset : offset + 24]
                )
                prop_array["time_type"] = punit_bytes[offset + 24 : offset + 25].decode(
                    "latin1"
                )
                prop_array["on_grid"] = punit_bytes[offset + 25 : offset + 26] == b"y"

                # Calculate net_len and read data
                net_len = char_len - 8 - 1 - MAX_STRING_LEN
                if prop_array["type"] == "d":
                    prop_array["length"] = net_len // 8
                    prop_array["data"] = np.frombuffer(
                        dym_prop_fp.read(net_len), dtype=np.float64
                    )
                elif prop_array["type"] == "i":
                    prop_array["length"] = net_len // 4
                    prop_array["data"] = np.frombuffer(
                        dym_prop_fp.read(net_len), dtype=np.int32
                    )
                else:
                    raise ValueError(f"Unknown data type: {prop_array['type']}")

                if np.isnan(prop_array["data"]).any():
                    print(f"Warning: NaN values found in property {name}")
                if np.isinf(prop_array["data"]).any():
                    print(f"Warning: Inf values found in property {name}")

                # Handle remapping if needed
                if remap and prop_array["on_grid"]:
                    if hasattr(self, "Moves") and self.Moves:
                        prop_array["data"] = self.re_map(
                            prop_array["data"],
                            self.NetSize(),
                            self.origin_net_size,
                            self.Moves,
                        )

                return True
        except FileNotFoundError:
            print(f"Error: File '{filename}' not found.")
            return False

    def read_well_data(self):
        """
        Read well data from the output file.

        """
        filename = f"{self.casename}_wstu.out"
        print(f"Reading WELL properties from {filename}")
        if not os.path.exists(filename):
            print(f'File "{filename}" does not exist!')
            return

        raw_data = {}
        current_time_data = None

        with open(filename, "r") as well_file:
            file_content = well_file.readlines()
            raw_data["XCoord"] = float(file_content[0].split()[1])
            raw_data["YCoord"] = float(file_content[1].split()[1])
            raw_data["TimeData"] = []
            for line in file_content[2:]:
                line = line.strip()
                if not line:
                    continue
                if line.startswith("Time"):
                    if current_time_data:
                        raw_data["TimeData"].append(current_time_data)
                    current_time_data = {"Time": float(line.split()[1]), "Wells": []}

                elif line.startswith("WELL"):
                    well_info = line.split()
                    well_id = int(well_info[1])
                    well_name = well_info[2]
                    well_data_num = int(well_info[3])
                    well_prop = well_info[4]

                    current_time_data["Wells"].append(
                        {
                            "ID": well_id,
                            "Name": well_name,
                            "DataNum": well_data_num,
                            "Prop": well_prop,
                            "Data": {},
                        }
                    )

                elif (
                    current_time_data
                    and "Wells" in current_time_data
                    and current_time_data["Wells"]
                ):
                    well_data = line.split()
                    prop = well_data[0]
                    unit = well_data[1]
                    real_data = well_data[2:]
                    current_time_data["Wells"][-1]["Data"][prop] = {
                        "Unit": unit,
                        "Data": real_data,
                    }

            if current_time_data:
                raw_data["TimeData"].append(current_time_data)

        print(f"Found {len(raw_data['TimeData'])} time steps in well data.")
        print(f"Found {len(raw_data['TimeData'][0]['Wells'])} wells in well data.")
        print(
            f"Found {len(raw_data['TimeData'][1]['Wells'][0]['Data'])} properties in well data."
        )

        # Processing well data
        self.well_data.XCoord = raw_data["XCoord"]
        self.well_data.YCoord = raw_data["YCoord"]
        self.well_data.timesteps = len(raw_data["TimeData"])
        self.well_data.well_num = len(raw_data["TimeData"][0]["Wells"])

        for well_data in raw_data["TimeData"][0]["Wells"]:
            well_info = {
                "ID": well_data["ID"],
                "Name": well_data["Name"],
                "DataNum": well_data["DataNum"],
                "Prop": well_data["Prop"],
                "Xcoord": np.array(
                    well_data["Data"]["XCOORD"]["Data"], dtype=np.float64
                ),
                "Ycoord": np.array(
                    well_data["Data"]["YCOORD"]["Data"], dtype=np.float64
                ),
                "Depth": np.array(well_data["Data"]["DEPTH"]["Data"], dtype=np.float64),
                "Stage": np.array(well_data["Data"]["STAGE"]["Data"], dtype=np.int64),
                "Outlet": np.array(well_data["Data"]["OUTLET"]["Data"], dtype=np.int64),
            }
            self.well_data.stat_data.append(well_info)
            self.well_data.points_num += well_data["DataNum"]

            well_data["Data"].pop("XCOORD")
            well_data["Data"].pop("YCOORD")
            well_data["Data"].pop("DEPTH")
            well_data["Data"].pop("STAGE")
            well_data["Data"].pop("OUTLET")

        for timestep_data in raw_data["TimeData"]:
            time_info = {
                "Time": timestep_data["Time"],
                "Data": [],
            }

            for well_data in timestep_data["Wells"]:
                single_data = {}
                single_data["ID"] = well_data["ID"]
                for prop in well_data["Data"]:
                    single_data[prop] = np.array(
                        well_data["Data"][prop]["Data"], dtype=np.float64
                    )
                time_info["Data"].append(single_data)

            self.well_data.dym_data.append(time_info)

    def scan(self, casename):
        """
        Scan the binary files for the specified case.

        :param casename: Case name.

        """
        self.clear()
        self.casename = casename

        # Read dynamic properties
        print(f"Reading dynamic properties for case {casename}")
        fcount = 0
        while True:
            filename = f"{self.casename}_stu{fcount if fcount != 0 else ''}.bin"
            if not os.path.exists(filename):
                if fcount == 0:
                    print(f'File "{filename}" does not exist!')
                break

            with open(filename, "rb") as fp:
                cur_time = -float("inf")
                while True:
                    current_position = fp.tell()
                    if not fp.read(1):
                        break
                    fp.seek(current_position, os.SEEK_SET)

                    tmp_prop = NameLoc()
                    tmp_prop.fcount = fcount
                    self.scan_bin_array(fp, tmp_prop)

                    if tmp_prop.time != cur_time or not self.dym_props:
                        step = len(self.dym_props)
                        for s in range(len(self.dym_props)):
                            if (
                                self.dym_props[s]
                                and self.dym_props[s][0].time == tmp_prop.time
                            ):
                                step = s
                                break

                        if step < len(self.dym_props):
                            self.dym_props[step].append(tmp_prop)
                        else:
                            self.dym_props.append([tmp_prop])

                        cur_time = tmp_prop.time
                    else:
                        self.dym_props[-1].append(tmp_prop)
            fcount += 1

        # Read static properties
        filename = f"{self.casename}_geom.bin"
        print(f"Reading static properties from {filename}")
        if os.path.exists(filename):
            with open(filename, "rb") as stat_prop_fp:
                while True:
                    current_position = stat_prop_fp.tell()
                    if not stat_prop_fp.read(1):
                        break
                    stat_prop_fp.seek(current_position, os.SEEK_SET)

                    tmp_prop = NameLoc()
                    self.scan_bin_array(stat_prop_fp, tmp_prop)

                    # seems like useless properties
                    if tmp_prop.name == "G2NMAP":
                        continue
                        gross_size = tmp_prop.length
                        stat_prop_fp.seek(
                            tmp_prop.offset + 17 + MAX_STRING_LEN, os.SEEK_SET
                        )
                        data_bytes = stat_prop_fp.read(gross_size * 4)
                        self.Gross2Net = np.frombuffer(data_bytes, dtype=np.int32)

                    if tmp_prop.name == "N2GMAP":
                        continue
                        net_size = tmp_prop.length
                        stat_prop_fp.seek(
                            tmp_prop.offset + 17 + MAX_STRING_LEN, os.SEEK_SET
                        )
                        data_bytes = stat_prop_fp.read(net_size * 4)
                        self.Net2Gross = np.frombuffer(data_bytes, dtype=np.int32)

                    self.stat_props.append(tmp_prop)
        else:
            print(f'File "{filename}" does not exist!')

        # Read well properties
        self.read_well_data()

        # Get GRIDINFO
        GRIDINFO_array = {}
        if self.get_static_data("GRIDINFO", GRIDINFO_array):
            if GRIDINFO_array["length"] == 32:
                self.GridInfo[:20] = GRIDINFO_array["data"][:20].astype(np.int32)
                self.AxisOffset = GRIDINFO_array["data"][20:23]
            else:
                print("Error: GRIDINFO length is not 32")
        else:
            print("Error: GRIDINFO data not found")

        # Consistency check
        if self.GridInfo[2] != len(self.Gross2Net):
            print("Error: G2NMAP length does not match GRIDINFO[2]")

        if self.GridInfo[8] != len(self.Net2Gross):
            print("Error: N2GMAP length does not match GRIDINFO[8]")

        self.origin_net_size = len(self.Net2Gross)
        self.origin_GridInfo = self.GridInfo.copy()
        self.origin_Gross2Net = self.Gross2Net.copy()

        self.print_info()

    def print_grid_info(self):
        print(
            f"Overall grid size = {self.get_gross_size()}, active grid = {self.get_net_size()}, number of time steps = {self.get_time_steps()}"
        )
        print(f"Axis_offset = {tuple(self.AxisOffset)}")
        print(
            f"NI = {self.get_grid_ni()}, NJ = {self.get_grid_nj()}, NK = {self.get_grid_nk()}, NLayer = {self.get_grid_nk() // self.get_porolvl()}"
        )
        print(
            f"PoroLvl = {self.get_porolvl()}, number of visible elements = length(NODE_LIST)/8 = {self.get_net_size()}"
        )

    def print_grid_nums(self):
        print(
            f"Base grid number = {self.get_base_grid_num()}, active base grid number = {self.get_active_base_grid_num()}"
        )
        print(
            f"Fine grid number = {self.get_fine_grid_num()}, active fine grid number = {self.get_active_fine_grid_num()}"
        )

    def print_static_props(self):
        print("Static properties:")
        for info in self.stat_props:
            print(
                f"    Name: {info.name}, length: {info.length}, unit: {info.unit}, is at cell center: {info.on_grid}"
            )

    def print_dynamic_props(self):
        print("Dynamic properties:")
        for dynamic_infos in self.dym_props:
            if dynamic_infos:
                print(
                    f"  Time: {dynamic_infos[0].time}, time type: {dynamic_infos[0].time_type}"
                )
            for info in dynamic_infos:
                print(
                    f"    Name: {info.name}, length: {info.length}, unit: {info.unit}, is at cell center: {info.on_grid}"
                )


if __name__ == "__main__":
    print("This is a reader class for binary files.")
